import cv2
import base64
import numpy as np
import torch
from flask import Flask, request, jsonify
from pyzbar.pyzbar import decode

app = Flask(__name__)

# 导入YOLOv5模型
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_boxes
from utils.dataloaders import letterbox

weights = 'runs/train/best.pt'  # 修改为你自己的模型路径
device = 'cpu'  # 使用的GPU设备号
imgsz = 640  # 输入图像尺寸

model = attempt_load(weights)
model.eval()


def detect_objects(image):
    # 对输入图像进行预处理，使其符合模型的要求
    img = letterbox(image, new_shape=imgsz)[0]

    # 转换为PyTorch的Tensor格式
    img = img[:, :, ::-1].transpose(2, 0, 1)
    img = np.ascontiguousarray(img)

    # 推断目标检测结果
    img = torch.from_numpy(img).to(device)
    img = img.float()
    img /= 255.0

    if img.ndimension() == 3:
        img = img.unsqueeze(0)

    # 进行目标检测
    pred = model(img)[0]

    # 后处理：非最大抑制
    pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.45)


    # 保存帧图像
    cv2.imwrite('guan.jpg', image)

    return pred


@app.route('/detect', methods=['GET'])
def start_detection():
    try:
        # 打开摄像头（0表示默认摄像头）
        cap = cv2.VideoCapture(0)

        ret, frame = cap.read()  # 读取摄像头帧

        # 执行目标检测
        detections = detect_objects(frame)

        # 调试输出，检查 detections 的值和长度
        print("Detections:", detections)
        print("Number of detections:", len(detections))

        # 处理检测结果
        if len(detections) > 0 and len(detections[0]) > 0:

            # 如果检测到目标，保存带有检测框的图像
            for det in detections[0]:
                x1, y1, x2, y2, conf, cls = det.tolist()
                x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.imwrite('detected_image.jpg', frame)

            # 使用pyzbar来识别二维码
            decoded_objects = decode(frame)

            # 获取第一个检测结果的置信度
            confidence = detections[0][0][4]

            # 如果置信度大于0.8，返回True；否则返回False
            if confidence > 0.8:
                # 如果识别到二维码，输出二维码内容
                if decoded_objects:
                    qr_code_content = decoded_objects[0].data.decode('utf-8')
                    return jsonify({"detected": True, "message": "识别到二维码并检测到内容", "qr_code_content": qr_code_content})
                else:
                    return jsonify({"detected": False, "message": "识别到二维码未检测到内容", "qr_code_content": None})
            else:
                return jsonify({"detected": False, "message": "识别到二维码但置信度小于0.8"})
        else:
            # 如果未检测到目标，返回False
            return jsonify({"detected": False, "message": "未识别到二维码"})
    except Exception as e:
        return str(e)


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
